"""
影刀数据库表结构修复脚本
"""
import sys
import os
import logging
from sqlalchemy import text

# 添加项目根目录到PATH
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from backend.database.database import engine, SessionLocal
from backend.wingknife.models import Base

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def check_column_exists(table_name, column_name):
    """检查表中是否存在特定列"""
    with engine.connect() as conn:
        query = text(f"PRAGMA table_info({table_name})")
        result = conn.execute(query)
        columns = [row[1] for row in result]
        return column_name in columns

def rename_column(table_name, old_column, new_column):
    """重命名表中的列 (SQLite不直接支持ALTER COLUMN，需要创建新表)"""
    with engine.connect() as conn:
        # 1. 检查表是否存在
        if not engine.dialect.has_table(conn, table_name):
            logger.info(f"表 {table_name} 不存在，无需重命名列")
            return False
            
        # 2. 获取表结构
        query = text(f"PRAGMA table_info({table_name})")
        result = conn.execute(query)
        columns = [(row[1], row[2]) for row in result]  # 列名和类型
        
        # 3. 检查列是否存在
        if old_column not in [col[0] for col in columns]:
            logger.info(f"列 {old_column} 在表 {table_name} 中不存在，无需重命名")
            return False
        
        # 4. 检查新列是否已存在
        if new_column in [col[0] for col in columns]:
            logger.info(f"列 {new_column} 在表 {table_name} 中已存在，无需重命名")
            return False
        
        # 5. 创建临时表
        temp_table = f"{table_name}_temp"
        create_temp_table = f"CREATE TABLE {temp_table} AS SELECT "
        
        # 构建SELECT部分
        select_parts = []
        for col_name, col_type in columns:
            if col_name == old_column:
                select_parts.append(f"{col_name} AS {new_column}")
            else:
                select_parts.append(col_name)
        
        create_temp_table += ", ".join(select_parts)
        create_temp_table += f" FROM {table_name}"
        
        # 6. 创建临时表并复制数据
        conn.execute(text(create_temp_table))
        
        # 7. 删除原表
        conn.execute(text(f"DROP TABLE {table_name}"))
        
        # 8. 重命名临时表
        conn.execute(text(f"ALTER TABLE {temp_table} RENAME TO {table_name}"))
        
        conn.commit()
        logger.info(f"已将表 {table_name} 中的列 {old_column} 重命名为 {new_column}")
        return True

def add_column(table_name, column_name, column_type, column_default=None):
    """向表中添加列"""
    with engine.connect() as conn:
        # 检查表是否存在
        if not engine.dialect.has_table(conn, table_name):
            logger.info(f"表 {table_name} 不存在，无需添加列")
            return False
            
        # 检查列是否已存在
        if check_column_exists(table_name, column_name):
            logger.info(f"列 {column_name} 在表 {table_name} 中已存在，无需添加")
            return False
        
        # 添加列
        alter_sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
        if column_default is not None:
            alter_sql += f" DEFAULT {column_default}"
        
        conn.execute(text(alter_sql))
        conn.commit()
        logger.info(f"已向表 {table_name} 添加列 {column_name}")
        return True

def fix_wingknife_schema():
    """修复影刀表结构"""
    logger.info("开始修复影刀表结构...")
    
    # 1. 创建表（如果不存在）
    Base.metadata.create_all(bind=engine)
    logger.info("已创建缺失的表")
    
    # 2. 检查并修复列名不匹配的问题
    try:
        # 检查processes表是否有config列，如果有则重命名为parameters
        if check_column_exists("wingknife_processes", "config") and not check_column_exists("wingknife_processes", "parameters"):
            rename_column("wingknife_processes", "config", "parameters")
        
        # 检查execution_records表是否有params列，如果有则重命名为parameters
        if check_column_exists("wingknife_execution_records", "params") and not check_column_exists("wingknife_execution_records", "parameters"):
            rename_column("wingknife_execution_records", "params", "parameters")
        
        # 添加缺失的列
        add_column("wingknife_processes", "output_dir", "VARCHAR(255)", "'uploads/wingknife/output/default'")
        add_column("wingknife_execution_records", "created_at", "DATETIME", "CURRENT_TIMESTAMP")
        add_column("wingknife_execution_records", "created_by", "VARCHAR(50)", "'system'")
        
        logger.info("表结构修复完成")
    except Exception as e:
        logger.error(f"修复表结构时出错: {str(e)}")
        return False
    
    return True

if __name__ == "__main__":
    success = fix_wingknife_schema()
    if success:
        logger.info("影刀表结构修复成功")
    else:
        logger.error("影刀表结构修复失败") 